# -*- coding: utf-8 -*-
import os
import cv2 as cv
import math
import numpy as np
import time


def dark_channel(normal_img, win_size):
    """
    求暗通道图
    :param normal_img: 归一化后的图像
    :param win_size: 最小值滤波的窗口大小
    :return: 暗通道图
    """
    b, g, r = cv.split(normal_img)
    # 求出BGR中数值最小者
    min_color = cv.min(cv.min(r, g), b)
    # 获取结构元素
    kernel = np.ones((2 * win_size + 1, 2 * win_size + 1))
    # 这个函数叫腐蚀 做的是win_size大小的模板运算 ,对应的是最小值滤波,即 黑色图像中的一块块的东西
    dark_img_f = cv.erode(min_color, kernel)
    return dark_img_f


def atm_light(normal_img, dark_img_f):
    """
    求全球大气光成分
    :param normal_img:归一化后的图像
    :param dark_img_f:暗通道图像
    :return:全球大气光成分
    """
    h, w = normal_img.shape[0:2]
    img_size = h * w
    num_px = int(max(math.floor(img_size / 1000), 1))
    dark_vec = dark_img_f.reshape(1, img_size)[0]
    img_vec = normal_img.reshape(img_size, 3)
    indices = dark_vec.argsort()
    indices = indices[::-1]
    atm_sum = np.zeros([1, 3])
    for ind in range(num_px):
        atm_sum = atm_sum + img_vec[indices[ind]]
    a = atm_sum / num_px
    a = a.reshape(1, 3)[0]
    return a


def guided_filter(gray_img, p, r, eps):
    """
    导向滤波
    :param gray_img: 灰度图像(归一化)
    :param p:估计的折射率
    :param r:r一般要大于最小值滤波半径的4倍(这里取60)4r+1
    :param eps:0.0001
    :return:优化后的折射率
    """
    mean_i = cv.boxFilter(gray_img, cv.CV_64F, (r, r))
    mean_p = cv.boxFilter(p, cv.CV_64F, (r, r))
    mean_ip = cv.boxFilter(gray_img * p, cv.CV_64F, (r, r))
    cov_ip = mean_ip - mean_i * mean_p

    mean_ii = cv.boxFilter(gray_img * gray_img, cv.CV_64F, (r, r))
    var_i = mean_ii - mean_i * mean_i

    a = cov_ip / (var_i + eps)
    b = mean_p - a * mean_i

    mean_a = cv.boxFilter(a, cv.CV_64F, (r, r))
    mean_b = cv.boxFilter(b, cv.CV_64F, (r, r))

    q = mean_a * gray_img + mean_b
    return q


def transmission_estimate(normal_img, a, win_size):
    """
    折射率估计
    :param normal_img:
    :param a:
    :param win_size:
    :return:
    """
    omega = 0.95
    im3 = np.empty(normal_img.shape, normal_img.dtype)

    for ind in range(0, 3):
        im3[:, :, ind] = normal_img[:, :, ind] / a[ind]

    transmission_e = 1 - omega * dark_channel(im3, win_size)
    return transmission_e


def transmission_refine(src_img_f, t_e, tx=0.1):
    """
    折射率优化
    :param src_img_f: 原始图像
    :param t_e: 估计的折射率
    :param tx: 折射率阈值
    :return: 优化后的折射率
    """
    gray1 = cv.cvtColor(src_img_f, cv.COLOR_BGR2GRAY)
    gray = np.float64(gray1) / 255
    t_f = guided_filter(gray, t_e, 60, 0.0001)
    # 设置折射率阈值，保护天空
    t_x = np.full(t_f.shape[0:2], tx, t_f.dtype)
    t_f = np.maximum(t_f, t_x)
    return t_f


def recover(normal_img, t_f, a):
    """
    :param normal_img: 归一化后的图像
    :param t_f: 折射率
    :param a: 全球大气光成分
    :return: 去雾后的图像矩阵
    """
    res = np.empty(normal_img.shape, normal_img.dtype)

    for i in range(0, 3):
        res[:, :, i] = (normal_img[:, :, i] - a[i]) / t_f + a[i]

    return res


def dark_channel_processing():
    my_path = ["/diskb/TZB/TZB1_royolo5/convertor/org/images/val",
               "/diskb/TZB/TZB1_royolo5/convertor/d/images/val"]
    img_list = os.listdir(my_path[0])
    print("文件中共 %d 张图片" % len(img_list))
    for count in range(len(img_list)):
        name = img_list[count]
        print(name)
        print("暗通道处理第 %d 张图片------>" % (count + 1), end='')
        start = time.time()
        os.chdir(my_path[0])
        # 读入图像
        src_img = cv.imread(img_list[count])
        # BGR图像归一化
        normalization_img = src_img.astype('float64') / 255
        # 转换为暗通道图像
        dark_img = dark_channel(normalization_img, 5)

        # 求出<<全球大气光成分>>A
        a = atm_light(normalization_img, dark_img)

        # 折射率（大气传递系数）估计
        te = transmission_estimate(normalization_img, a, 4)

        # 折射率（大气传递系数）优化
        t = transmission_refine(src_img, te, 0.1)

        # 图像去雾
        result_img = recover(normalization_img, t, a) * 255

        # 保存图片
        os.chdir(my_path[1])
        cv.imwrite(img_list[count], result_img)
        print("第 %d 张图片暗通道处理完成" % (count + 1), end='>>>')
        print('处理用时：', time.time() - start)


dark_channel_processing()
